# 引入numpy, plt, io, cv
import numpy as np
import matplotlib.pyplot as plt
from skimage import io
import cv2 as cv
from skimage.filters import threshold_otsu  # 引入skimage的otsu函数

# 1. 读取图像
img = io.imread('cuteCat.jpg')  # 读取RGB有标签猫图：cuteCat.jpg 【注意大小写】
g = cv.imread('cuteCat.jpg', 0)  # 读取猫图 cuteCat.jpg 灰度图
h, w, _ = img.shape  # 获取猫图 cuteCat.jpg 的高和宽

# 2. 制作掩模
r = img[:, :, 0]  # 读取图像红色通道内容
k = 1.5  # 严格系数设置为1.5
r_mark = r[0:120, 280:]  # 取r中右上角，只包含标记的区域高度0到120，280到670
thresh = threshold_otsu(r_mark)  # 用otsu算法取区域中的阈值
mask = (r > g * k) & (r > thresh)  # 使 全局阈值 和 局部阈值 共同施加作用
# mask为基于0和1的二值图像
# 因为OpenCV只接受ubyte类型的图像，所以要乘以255，把0和1的二值图像编程0和255的np.ubyte图像
mask = (mask * 255).astype(np.ubyte)

# 注意 erode 和 dilate 都需要一个全都是1的小邻域窗口
# 窗口大小为3~5较合适, 数据类型为np.ubyte
kernel = np.ones((3, 3), np.ubyte)
mask = cv.erode(mask, kernel, iterations=1)  # 1次erode
mask = cv.dilate(mask, kernel, iterations=2)  # 2次dilate

# 3. 使用NS和TELEA进行填充
res1 = cv.inpaint(src=img, inpaintMask=mask, inpaintRadius=5, flags=cv.INPAINT_NS)  # NS填充
res2 = cv.inpaint(src=img, inpaintMask=mask, inpaintRadius=5, flags=cv.INPAINT_TELEA)  # TELEA填充

fig, ax = plt.subplots(2, 2, figsize=(10, 7))
ax[0, 0].imshow(img)
ax[0, 0].set_title("original", size=20)
ax[0, 1].imshow(mask, cmap='gray')
ax[0, 1].set_title("mask: r>g*1.5 and r>" + str(thresh), size=20)
ax[1, 0].imshow(res1)
ax[1, 0].set_title("INPAINT_NS", size=20)
ax[1, 1].imshow(res2)
ax[1, 1].set_title("INPAINT_TELEA", size=20)

plt.tight_layout()
plt.savefig('out.png')
plt.show()
